import torchvision
from torch import nn
# from torch.utils.tensorboard import SummaryWriter

from common.logging_format import *
from model_ import *
logger = Logger("train_log.log", "train_log")
train_dataset = torchvision.datasets.CIFAR10(root='../data/', train=True, download=False,
                                             transform=torchvision.transforms.ToTensor())

test_dataset = torchvision.datasets.CIFAR10(root='../data/', train=False, download=False,
                                            transform=torchvision.transforms.ToTensor())
logger.info("训练集数据长度:{}".format(len(train_dataset)))
logger.info("测试集数据长度:{}".format(len(test_dataset)))

train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=64)
test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=64)

"""
    有GPU的话，需要把：模型、数据（训练+测试）、损失函数、优化器 都放到cuda中
    方式一: 在每个需要gpu计算的地方加上代码：
    if torch.cuda.is_available():
        my_net = my_net.cuda()
    方式二: device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
            my_net = my_net.to(device)
    显卡id以0开始，因此只有一块显卡时，cuda:0和cuda含义相同
    方式二明显要优于方式一，因此一般选用方式二
"""
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
my_net = MyNet()
loss_fn = nn.CrossEntropyLoss()
my_net.to(device)
loss_fn = loss_fn.to(device)

optimizer = torch.optim.SGD(my_net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)

epochs = 5

total_test_loss = 0

total_test_step = 0

# writer = SummaryWriter("../logs")

total_train_step = 0
for epoch in range(epochs):
    logger.info('--------------------第{}轮训练--------------------'.format(epoch + 1))

    for images, labels in train_dataloader:
        images, labels = images.to(device), labels.to(device)
        output = my_net(images)
        loss = loss_fn(output, labels)
        output = output.to(device)
        loss = loss.to(device)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if total_train_step % 100 == 0:
            logger.info('Epoch [{}/{}], batch [{:4d}] Loss: {:.4f}'.format(epoch + 1, epochs, total_train_step, loss.item()))
            # writer.add_scalar('train_loss', loss.item(), total_train_step)
        total_train_step += 1
    total_test_loss = 0
    total_accuracy = 0
    with torch.no_grad():
        for data, target in test_dataloader:
            output = my_net(data)
            loss = loss_fn(output, target)
            output = output.to(device)
            loss = loss.to(device)
            total_test_loss += loss.item()
            # 计算正确率
            accuracy = (output.argmax(1) == target).sum()
            total_accuracy += accuracy
    logger.info('Test Loss: {:.4f}'.format(total_test_loss / len(test_dataset)))
    logger.info('Total accuracy: {:.4f}'.format(total_accuracy / len(test_dataset)))
    # writer.add_scalar('test_loss', total_test_loss, total_test_step)
    # writer.add_scalar('test_accuracy', total_accuracy, total_test_step)
    total_test_step += 1

    torch.save(my_net, 'model_{}.pth'.format(epoch))
# writer.close()
